import numpy as np
import cv2
import cv2.aruco as aruco
import os

class CalibrateRotation(object):
    def __init__(self, data_path, camera_param_path):

        self.data_path = data_path

        # 图像路径列表
        img_names = os.listdir(data_path)
        self.img_paths = [os.path.join(data_path, name) for name in img_names if name.endswith('.bmp')]

        # 相机内参矩阵和畸变系数
        self.load_camera_param(camera_param_path)

    def load_mark_pixels(self):
        pixels = []
        for img_path in self.img_paths:
            image = cv2.imread(img_path)
            newimage = self.undistort_image(image)
            mark_pixel = self.get_one_mark_pixel(newimage)
            pixels.append(mark_pixel)
        pixels = np.asarray(pixels)

        return pixels
    
    def get_one_mark_pixel(self, image):
        image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        # 取ArUco码的四个角点作为特征
        aruco_dict = aruco.getPredefinedDictionary(aruco.DICT_7X7_1000)
        parameters = aruco.DetectorParameters()
        detector = aruco.ArucoDetector(aruco_dict, parameters)
        corners, ids, rejected = detector.detectMarkers(image_gray)

        # print(corners[0][0,0,:])
        cv2.circle(image, (int(corners[0][0,0,0]), int(corners[0][0,0,1])), 2, (255,0,0), 5)
        cv2.namedWindow("img",0)
        cv2.resizeWindow("img", 640, 480)
        cv2.imshow("img", image)
        cv2.waitKey(0)

        temp = corners[0][0,:,:].tolist()
        if len(temp) != 4:
            raise ValueError("Detected marker corners do not contain exactly four points.")
        
        # 方法一：根据四个角点的坐标取中点坐标
        # center_point = self.get_center_point(temp)
        # 方法二：取二维码左上角坐标
        center_point = temp[0]

        return center_point
    
    def get_center_point(self, points):
        p1, p2, p3, p4 = points[0], points[1], points[2], points[3]
        # p1, p3
        a1, b1, c1 = p3[1]-p1[1], p1[0]-p3[0], p3[0]*p1[1] - p1[0]*p3[1]
        # p2, p4
        a2, b2, c2 = p4[1]-p2[1], p2[0]-p4[0], p4[0]*p2[1] - p2[0]*p4[1]
        # 求交点
        assert (a1*b2 - a2*b1) != 0, "两条直线平行，无法求交点！"
        x = (b1*c2 - b2*c1) / (a1*b2 - a2*b1)
        y = (a1*c2 - a2*c1) / (a2*b1 - a1*b2)
        # print("交点坐标: ", x, y)
        return (x, y)
    
    def fit_circle(self):
        """
        多点拟合圆
        输入：多个点的坐标 [(x1,y1), (x2,y2), (x3,y3)]
        输出：圆心坐标 (h,k) 和半径 r
        """
        points = self.load_mark_pixels()
        
        # # 构建矩阵 A 和向量 b
        A = np.column_stack((points, np.ones(points.shape[0])))


        b = -np.sum(points**2, axis=1).reshape(-1, 1)
        
        # 解线性方程组 Ax = b
        x = np.linalg.lstsq(A, b)[0]
        D, E, F = x.flatten()
        # 计算圆心和半径
        x_circle = -D/2
        y_circle = -E/2
        r = np.sqrt(x_circle**2 + y_circle**2 - F) / 2
        
        return (x_circle, y_circle), r
    
    def undistort_image(self, img):
        h, w = img.shape[:2]
        # 计算去畸变和裁剪的优化矩阵
        newcameramtx, roi = cv2.getOptimalNewCameraMatrix(self.cameraMatrix, self.distCoeffs, (w, h), 1, (w, h))

        # 去畸变
        dst = cv2.undistort(img, self.cameraMatrix, self.distCoeffs, None, newcameramtx)
        return dst

    def load_camera_param(self, camera_param_path):
        camera_param_cvfile = cv2.FileStorage(camera_param_path, cv2.FILE_STORAGE_READ)
        self.cameraMatrix = camera_param_cvfile.getNode("cameraMatrix").mat()
        self.distCoeffs = camera_param_cvfile.getNode("distCoeffs").mat()


if __name__ == "__main__":
    data_path = "D:/VSCodeProjects/calibration/data/roation202506131249"
    camera_param_path = "D:/VSCodeProjects/calibration/data/camera_params.yaml"
    caliroation = CalibrateRotation(data_path, camera_param_path)

    circle_point, r = caliroation.fit_circle()
    print("圆心坐标: ", circle_point)
    print("半径: ", r)

    # 创建一个空白图像
    # image = cv2.imread(os.path.join(data_path, "2.bmp"))

    # # 绘制轮廓
    # # cv2.polylines(image, [[[591,1172]]], isClosed=True, color=(0, 255, 0), thickness=2)

    # # 绘制圆心
    # cv2.circle(image, (1138, 2742), 3, (0, 0, 255), -1)

    # # 绘制圆
    # cv2.circle(image, (1138, 2742), int(r), (255, 0, 0), 2)

    # # 显示图像
    # # cv2.imshow("Contour and Fitted Circle", image)
    # # cv2.waitKey(0)
    # # cv2.destroyAllWindows()
    # cv2.imwrite("contour_and_fitted_circle.png", image)

